import cv2
import time
from ultralytics import YOLO

# 1. 地址
rtsp_url = "rtsp://admin:123456Wy%21@192.168.0.64:554/Streaming/Channels/102"

# 2. 打开流
cap = cv2.VideoCapture(rtsp_url, cv2.CAP_FFMPEG)
if not cap.isOpened():
    raise IOError("无法打开 RTSP 流")

# 3. 只加载一次模型
model = YOLO("yolov8n.pt")

# 4. 读帧 + 推理
while True:
    ret, frame = cap.read()
    if not ret:
        time.sleep(0.01)
        continue

    # 5. 推理（用 GPU 则自动走 CUDA）
    results = model(frame, verbose=False)
    annotated = results[0].plot()

    # 6. 显示 **画框后的图**
    cv2.imshow("hik", annotated)

    # 7. 按 q 或 ESC 退出
    if cv2.waitKey(1) & 0xFF in (27, ord('q')):
        break

cap.release()
cv2.destroyAllWindows()